# encoding:utf/8
from mmdet.apis import inference_detector, init_detector
import json
import os
from tqdm import tqdm
import cv2
import numpy as np


def get_iou(bbox1, bbox2):
    """

    Args:
        bbox1: the coor of the first bbox
        bbox2: the coor of the second bbox

    Returns:

    """
    area1 = (bbox1[2] - bbox1[0]) * (bbox1[3] - bbox1[1])
    area2 = (bbox2[2] - bbox2[0]) * (bbox2[3] - bbox2[1])
    left_colum_max = max(bbox1[0], bbox2[0])
    right_colum_min = min(bbox1[2], bbox2[2])
    up_row_max = max(bbox1[1], bbox2[1])
    down_row_min = min(bbox1[3], bbox2[3])
    if(left_colum_max >= right_colum_min or down_row_min <= up_row_max):
        return 0
    else:
        area_cross = (down_row_min - up_row_max)*(right_colum_min - left_colum_max)
        return area_cross/(area1+area2-area_cross)


def nms(result, thre=0.5):
    """
    Args:
        result: the result of the same categary
    Returns:
        bbox after result
    """
    scores = result[:, 4]
    x1 = result[:, 0]
    y1 = result[:, 1]
    x2 = result[:, 2]
    y2 = result[:, 3]
    areas = (y2-y1+1)*(x2-x1+1)
    keep = []
    index = scores.argsort()[::-1]
    while(index.size > 0):
        i = index[0]
        keep.append(i)
        x11 = np.maximum(x1[i], x1[index[1:]])
        y11 = np.maximum(y1[i], y1[index[1:]])
        x22 = np.minimum(x2[i], x2[index[1:]])
        y22 = np.minimum(y2[i], y2[index[1:]])
        w = np.maximum(0, x22-x11+1)
        h = np.maximum(0, y22-y11+1)
        overlaps = w*h
        ious = overlaps/(areas[i]+areas[index[1:]]-overlaps)
        idx = np.where(ious <= thre)[0]
        index = index[idx+1]
    return keep


def result_from_dir():
    index = {1: 1, 2: 2, 3: 3, 4: 4, 5: 5, 6: 6}
    # build the model from a config file and a checkpoint file
    model = init_detector(config_path, model_path, device='cuda:8')
    pics = os.listdir(pic_path)
    results = []
    num = 0
    result_ = []
    for im in tqdm(pics):
        num += 1
        img = os.path.join(pic_path, im)
        im = cv2.imread(img)
        H,W,C = im.shape
        im1 = im[:H//2, :W//2, :]
        im2 = im[:H//2, W//2:, :]
        im3 = im[H//2:, :W//2, :]
        im4 = im[H//2:, W//2:, :]
        im5 = cv2.resize(im, (W//2, H//2))
        img = [im1, im2, im3, im4, im5]
        for i in range(len(img)):
            res = inference_detector(model, img[i])
            if(i == 1):
                for r in res:
                    if(r.shape[0] > 0):
                        r[0][0] += W//2
                        r[0][2] += W//2
            elif(i == 2):
                for r in res:
                    if (r.shape[0] > 0):
                        r[0][1] += H//2
                        r[0][3] += H//2
            elif(i == 3):
                for r in res:
                    if (r.shape[0] > 0):
                        r[0][0] += W // 2
                        r[0][2] += W // 2
                        r[0][1] += H // 2
                        r[0][3] += H // 2
            elif(i == 4):
                for r in res:
                    if (r.shape[0] > 0):
                        r[0][0] *= 2
                        r[0][2] *= 2
                        r[0][1] *= 2
                        r[0][3] *= 2
            result_.append(res)
        result1 = []
        result2 = []
        result3 = []
        result4 = []
        result5 = []
        result6 = []
        for r in result_:
            result1.append(r[0])
            result2.append(r[1])
            result3.append(r[2])
            result4.append(r[3])
            result5.append(r[4])
            result6.append(r[5])
        result1 = np.vstack(result1)
        result2 = np.vstack(result2)
        result3 = np.vstack(result3)
        result4 = np.vstack(result4)
        result5 = np.vstack(result5)
        result6 = np.vstack(result6)
        # NMS
        result = []
        result_ = list((result1, result2, result3, result4, result5, result6))
        for i in result_:
            if(i.shape[0] > 1):
                keep_ids = nms(i, 0.5)
                result.append(i[keep_ids])
            else:
                result.append(i)

        # 从1开始
        for i, boxes in enumerate(result, 1):
            if len(boxes):
                defect_label = index[i]
                for box in boxes:
                    d = {}
                    d["name"] = im
                    d['category'] = defect_label
                    d['bbox'] = [round(float(i), 2) for i in box[0:4]]
                    d['score'] = float(box[4])
                    results.append(d)

    with open(json_out_path, 'w') as fp:
        json.dump(results, fp)


if __name__ == "__main__":
    model_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/latest.pth'
    config_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/one2four_add_mjx.py'
    json_out_path = '/data/lzy/work_dir/c_rcnn_r50_fpn_2x_coco_RCCP_640_1280_add_ratio_add_dcn_add_size4_one2four_addmjx_no_dataaug/result_bak.json'
    pic_path = '/data/lzy/tile_round1_testA_20201231/testA_imgs/'
    result_from_dir()


